import torch
import argparse
from PIL import Image   # in data
from os.path import join
import os
from torchvision import transforms  # in data
from anomaly.models import PatchCore
from torch.utils.data import DataLoader
from anomaly.datasets import MVTecDataset, mean_train, std_train



def get_args():
    parser = argparse.ArgumentParser(description='ANOMALYDETECTIONPatchCore')
    parser.add_argument('--phase', choices=['train','test'], default='test')
    parser.add_argument('--dataset_path', default='../../_DATASET/mvtec')
    parser.add_argument('--category', default='bottle')
    parser.add_argument('--num_epochs', default=1, type=int)
    parser.add_argument('--batch_size', default=32)
    parser.add_argument('--load_size', default=256) # 256
    parser.add_argument('--input_size', default=224)
    parser.add_argument('--coreset_sampling_ratio', default=0.001)  
    parser.add_argument("--weights_dir", type=str, default='../../_Weights/anomaly_lab/patchcore/weights', help="Define where to save model checkpoints.")
    parser.add_argument('--result_dir', default='../../_Weights/anomaly_lab/patchcore/result',help='Experiment name (defult None).') # 性能指标文件, 可视化在root/sample
    parser.add_argument('--save_anomaly_map', default=True)
    parser.add_argument('--n_neighbors', type=int, default=9)
    args = parser.parse_args()

    return args


if __name__ == '__main__':

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    args = get_args()

    if not os.path.exists(args.weights_dir):
        os.makedirs(args.weights_dir)

    if not os.path.exists(args.result_dir):
        os.makedirs(args.result_dir)
        
    
    # for patchcore data
    data_transforms = transforms.Compose([
                            transforms.Resize((args.load_size, args.load_size), Image.ANTIALIAS),
                            transforms.ToTensor(),
                            transforms.CenterCrop(args.input_size),
                            transforms.Normalize(mean=mean_train,
                                                std=std_train)])
    gt_transforms = transforms.Compose([
                    transforms.Resize((args.load_size, args.load_size)),
                    transforms.ToTensor(),
                    transforms.CenterCrop(args.input_size)])

    inv_normalize = transforms.Normalize(mean=[-0.485/0.229, -0.456/0.224, -0.406/0.255], std=[1/0.229, 1/0.224, 1/0.255])

    model = PatchCore(args)

    train_datasets = MVTecDataset(root=join(args.dataset_path,args.category), transform=data_transforms, gt_transform=gt_transforms, phase='train')
    train_loader = DataLoader(train_datasets, batch_size=args.batch_size, shuffle=True, num_workers=0)

    test_datasets = MVTecDataset(root=join(args.dataset_path,args.category), transform=data_transforms, gt_transform=gt_transforms, phase='test')
    test_loader = DataLoader(test_datasets, batch_size=1, shuffle=False, num_workers=0)

    
    if args.phase == 'train':
        model.train(train_loader)
        model.train_after()
        # model.test_after(test_loader)
        # model.evaluate()
    elif args.phase == 'test':
        model.test_after(test_loader)
        model.evaluate()